import numpy as np
import seaborn as sns
from matplotlib import pyplot as plt

from model import *
from model.data_processing.data_loaders import EndJointData
from model.data_processing.jax_dataloader import NumpyLoader
from model.transforms import denormalize_joint, normalize_end

data = EndJointData("data/json_data/can0_L20Full.json")
data_loader = NumpyLoader(
    data,
    batch_size=512,
    shuffle=True,
    num_workers=10,
    pin_memory=False,
)


###########  weighted!
# model = MLPCtrlModel(in_dim=3, out_dim=6, model_file="nn_end_joint_v0.ckp")
# model = MLPCtrlModel(in_dim=3, out_dim=6, model_file="nn_end_joint_old.ckp")

model = DiffusionIKModel(
    in_dim=3,
    out_dim=6,
    model_file="diffusion_end_joint_all_normed.ckp",
    dropout_rate=0.1,
    num_act_samples=500,
    temperature=0.1,
    fk_model=MlpPartialIKModel(in_dim=6, out_dim=3, model_file="all_normed_FK.ckp"),
)


# model.fk_model = MLPCtrlModel(in_dim=6, out_dim=3, model_file="mlp_normed_FK.ckp").model


x = data.ends.numpy()[:, :3]
y = data.joints.numpy()
pred = model(normalize_end(x[:1000]))

error = np.sqrt((denormalize_joint(pred) - y[:1000]) ** 2)
mse = error.mean(axis=0)


for i, e in enumerate(error.T):
    sns.kdeplot(e, label=f"joint{i}={mse[i]:.3f}")

plt.xlim(0, 10)
plt.legend()
plt.show()
